import os
from math import ceil

import pandas as pd
import torch
from gymnasium.wrappers import FlattenObservation
from sb3_contrib.qrdqn import QRDQN
from stable_baselines3.common.vec_env import SubprocVecEnv

from MaskableQRDQNMultiInputPolicy import MaskableQRDQNMultiInputPolicy, MaskingReplayBuffer
from TransformerFeatureExtractor import TransformerFeatureExtractor
from trade_account import TradeAccount
from trade_env import TradeEnv

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

def check_timestamp_consistency(df: pd.DataFrame, time_col: str = 'timestamp') -> None:
    if time_col not in df.columns:
        raise ValueError(f"列 '{time_col}' 不存在")
    if not pd.api.types.is_datetime64_any_dtype(df[time_col]):
        raise TypeError(f"列 '{time_col}' 不是 datetime 类型")

    df_sorted = df.sort_values(by=time_col).reset_index(drop=True)
    time_diffs = df_sorted[time_col].diff().dropna()
    expected_diff = time_diffs.iloc[0]
    inconsistent = time_diffs != expected_diff

    if inconsistent.any():
        raise ValueError(f"❌ 时间间隔不一致，共 {inconsistent.sum()} 条异常记录")
    print("✅ 时间戳连贯，间隔一致")

def make_single_env(df, tech_indicator_list, seq_len):
    def _init():
        env = TradeEnv(df=df.copy(), tech_indicator_list=tech_indicator_list, seq_len=seq_len)
        return env
    return _init

def make_vec_env(df, tech_indicator_list, num_envs, seq_len):
    df = df.reset_index(drop=True)
    chunk_size = ceil(len(df) / num_envs)

    env_fns = []
    for i in range(num_envs):
        start = i * chunk_size
        end = min((i + 1) * chunk_size, len(df))
        sub_df = df.iloc[start:end].reset_index(drop=True)
        print(f"环境 {i+1} 分配数据 {start}~{end}（共 {len(sub_df)} 条）")
        env_fns.append(make_single_env(sub_df, tech_indicator_list, seq_len))

    return SubprocVecEnv(env_fns)

class TradeAgent:
    def __init__(self, df, model_kwargs, tech_indicator_list, num_envs=1, device=DEVICE):
        df = df.copy()
        df['timestamp'] = pd.to_datetime(df['timestamp'])
        check_timestamp_consistency(df)

        # 更新特征提取器参数
        policy_kwargs = model_kwargs.get('policy_kwargs', {})
        fe_kwargs = policy_kwargs.get('features_extractor_kwargs', {})
        fe_kwargs['feature_dim'] = len(tech_indicator_list) + 5 + len(TradeAccount().get_account_state())
        fe_kwargs['seq_len'] = fe_kwargs.get('seq_len', 60)

        # 确保特征提取器类正确设置
        policy_kwargs['features_extractor_class'] = TransformerFeatureExtractor
        policy_kwargs['features_extractor_kwargs'] = fe_kwargs
        policy_kwargs['normalize_images'] = True
        model_kwargs['policy_kwargs'] = policy_kwargs
        model_kwargs['device'] = device
        self.model_kwargs = model_kwargs
        self.env = make_vec_env(df, tech_indicator_list, num_envs, fe_kwargs['seq_len'])
        self.device = device
        self.df = df
    def get_model(self):
        print(f"model_kwargs = {self.model_kwargs}")
        model = QRDQN(
            verbose  = 1,
            policy = MaskableQRDQNMultiInputPolicy,
            replay_buffer_class=MaskingReplayBuffer,
            env = self.env,
            **self.model_kwargs
        )
        return model

    @staticmethod
    def save_model(path, model):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        model.save(path)
        print(f"[模型已保存至 {path}]")

    @staticmethod
    def load_model(path, device=DEVICE, custom_objects=None):
        if not os.path.exists(path):
            print(f"模型文件不存在: {path}")
            return None

        model = QRDQN.load(path, device=device, custom_objects=custom_objects)
        print(f"[模型已加载 {path}]")
        return model

    def train_model(
            self,
            model_save_path,
            model_load_path=None,
            custom_objects=None,
            step_multiples=1,
    ):
        model = None
        if model_load_path is not None:
            print(f"加载已有模型 {model_load_path} 进行增量训练")
            model = self.load_model(path=model_load_path, device=self.device, custom_objects=custom_objects)
            model.set_env(self.env)

        if model is None:
            print(f"模型加载失败或未提供，创建新模型")
            model = self.get_model()

        model.learn(
            total_timesteps=int(len(self.df) * step_multiples),
            progress_bar=True,
            reset_num_timesteps=False
        )

        self.save_model(model_save_path, model)
        print(f"✅ 训练完成，模型保存至 {model_save_path}")